package 左哥算法.ch18数学知识;

import org.junit.Test;

import java.util.Arrays;
import java.util.Scanner;

import static org.junit.Assert.assertArrayEquals;

public class 快速幂 {

    @Test
    public void test1() {
        System.out.println(pow(3, 3));
        System.out.println(pow(3, 1));
        System.out.println(pow(3, 0));
    }

    /**
     * 快速幂
     *
     * @param n
     * @param k
     * @return
     */
    public long pow(int n, int k) {
        long res = 1;
        long temp = n;
        while (k > 0) {
            if (k % 2 != 0) {
                res *= temp;
            }
            temp *= temp;
            k /= 2;
        }
        return res;
    }

    @Test
    public void testQuickPowCase1() {
        long[][] matrix = {{1, 2}, {3, 4}};
        long[][] result1 = new long[][]{{37, 54}, {81, 118}}; // 按照快速幂计算，矩阵^3的结果
        int k1 = 3;

        long[][] output1 = quickPow(matrix, k1);
        for (long[] longs : output1) {
            System.out.println(Arrays.toString(longs));
        }
        assertArrayEquals(result1, output1);
    }

    @Test
    public void testQuickPowCase2() {
        long[][] matrix = {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}};
        long[][] result2 = matrix; // 单位矩阵的任何次方都是单位矩阵
        int k2 = 10;


        long[][] output2 = quickPow(matrix, k2);
        for (long[] longs : output2) {
            System.out.println(Arrays.toString(longs));
        }
        assertArrayEquals(result2, output2);
    }

    @Test
    public void testQuickPowCase3() {
        long[][] matrix = {{2, 0}, {1, 2}};
        long[][] result3 = new long[][]{{4, 0}, {4, 4}}; // 矩阵^2的结果
        int k3 = 2;


        long[][] output3 = quickPow(matrix, k3);
        for (long[] longs : output3) {
            System.out.println(Arrays.toString(longs));
        }
        assertArrayEquals(result3, output3);
    }

    /**
     * 矩阵快速幂
     */
    public long[][] quickPow(long[][] matrix, int k) {
        if (matrix.length != matrix[0].length) return null;
        int n = matrix.length;
        long[][] res = new long[n][n];
        // 初始化单位矩阵
        for (int i = 0; i < res.length; i++) {
            res[i][i] = 1;
        }
        long[][] temp = matrix;
        while (k > 0) {
            if (k % 2 != 0) {
                res = multiply(res, temp);
            }
            temp = multiply(temp, temp);
            k /= 2;
        }
        return res;
    }

    @Test
    public void testMultiply2() {
        long[][] a = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};
        long[][] b = {{9, 8, 7}, {6, 5, 4}, {3, 2, 1}};

        long[][] result = multiply(a, b);

        long[][] expected = {{30, 24, 18}, {84, 69, 54}, {138, 114, 90}};
        assertArrayEquals(expected, result);
    }

    @Test
    public void testMultiply3() {
        long[][] a = {{1, 2}, {3, 4}};
        long[][] b = {{7, 10}, {15, 22}};

        long[][] result = multiply(a, b);
        for (long[] longs : result) {
            System.out.println(Arrays.toString(longs));
        }
    }

    /**
     * 实现两个矩阵相乘的函数
     *
     * @param a 第一个矩阵，行数为m，列数为n
     * @param b 第二个矩阵，行数为n，列数为p
     * @return 乘积矩阵，行数为m，列数为p
     */
    public long[][] multiply(long[][] a, long[][] b) {
        // 创建结果矩阵，行数为a的行数，列数为b的列数
        long[][] res = new long[a.length][b[0].length];
        for (int i = 0; i < res.length; i++) {
            for (int j = 0; j < res[i].length; j++) {
                for (int k = 0; k < b.length; k++) {
                    // 计算乘积矩阵的每个元素
                    res[i][j] += a[i][k] * b[k][j];
                }
            }
        }
        return res;
    }

    public long fibonacci2(int n) {
        if (n < 2) return 1;
        long[][] base = {
                {1, 1},
                {1, 0}
        };
        long[][] first = {{1, 1}}; //f2和f1
        long[][] longs = quickPow(base, n - 1);
        long[][] res = multiply(first, longs);
        return res[0][0];
    }
}
